package com.hzlj.position.locate.config;

import cn.hutool.core.date.DateUtil;
import cn.hutool.core.util.ObjectUtil;
import com.fz.common.base.core.date.DateUtils;
import com.fz.common.base.exception.PubException;
import com.google.common.collect.Lists;
import com.google.common.collect.Range;
import com.google.common.collect.Sets;
import lombok.extern.slf4j.Slf4j;
import org.apache.shardingsphere.sharding.api.sharding.complex.ComplexKeysShardingAlgorithm;
import org.apache.shardingsphere.sharding.api.sharding.complex.ComplexKeysShardingValue;
import org.springframework.stereotype.Component;

import java.sql.Timestamp;
import java.time.LocalDate;
import java.util.Collection;
import java.util.Date;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * 自定义分表
 *
 * @author lifh
 **/
@Slf4j
@Component("reportInitiativeShardingAlgorithm")
public class ReportInitiativeShardingAlgorithm implements ComplexKeysShardingAlgorithm<Date> {

    @Override
    public Collection<String> doSharding(Collection<String> collection, ComplexKeysShardingValue<Date> complexKeysShardingValue) {
        String table = complexKeysShardingValue.getLogicTableName();
        //=,in的情况
        Collection<Date> belongDays = complexKeysShardingValue.getColumnNameAndShardingValuesMap().get("belong_day");
        if (ObjectUtil.isNotEmpty(belongDays)) {
            return tables(table, getQuarters(belongDays));
        }
        //<>的情况
        Range<?> belongDaysRange = complexKeysShardingValue.getColumnNameAndRangeValuesMap().get("belong_day");

        Date startTime = parseBelongDay(belongDaysRange.lowerEndpoint());
        Date endTime = parseBelongDay(belongDaysRange.upperEndpoint());
        if (endTime.compareTo(startTime) < 0) {
            throw new PubException("结束日期不能小于开始日期");
        }
        return tables(table, getQuarters(startTime, endTime));
    }


    private Collection<String> tables(String table, Set<String> quarters) {
        return quarters.stream()
                .map(e -> table + "_" + e)
                .collect(Collectors.toList());
    }

    private Set<String> getQuarters(Date startTime, Date endTime) {
        List<Date> dates = Lists.newArrayList();
        LocalDate start = DateUtils.localDate(DateUtils.monthStart(startTime));
        LocalDate end = DateUtils.localDate(DateUtils.monthStart(endTime));
        // 初始化
        while (start.compareTo(end) <= 0) {
            dates.add(DateUtils.date(start));
            start = start.plusMonths(1);
        }
        return getQuarters(dates);
    }


    private Set<String> getQuarters(Collection<?> dates) {
        Collection<Date> parseDates = dates.stream().map(this::parseBelongDay).collect(Collectors.toList());
        Set<String> quarters = Sets.newHashSet();
        for (Date date : parseDates) {
            quarters.add(
                    DateUtil.year(date) + "_" + DateUtil.quarter(date)
            );
        }
        return quarters;
    }


    private Date parseBelongDay(Object belongDay) {
        if (ObjectUtil.isEmpty(belongDay)) {
            return DateUtils.start();
        }
        Date day = null;
        if (belongDay instanceof String || belongDay instanceof Integer) {
            day = DateUtil.parse(belongDay.toString());
        } else if (belongDay instanceof Timestamp) {
            day = new Date(((Timestamp) belongDay).getTime());
        } else if (belongDay instanceof Date) {
            day = (Date) belongDay;
        }
        if (day != null) {
            return DateUtil.beginOfDay(day);
        }

        throw new PubException("日期格式不能为空");
    }

    @Override
    public void init() {
    }

    @Override
    public String getType() {
        return null;
    }
}
